#include "file_manager.h"

#include <unistd.h>

#include <filesystem>

#include "fitsio.h"

#include "dwarf_utils.h"
#include "elog.h"
#include "elog_file.h"
#include "timer_utils.h"

using namespace cv;
using namespace std;
namespace fs = std::filesystem;

const ImgInfo FileManager::imgInfo4K_ = { Size(3840, 2160), 3840 * 2160 * sizeof(ushort), 3840 * 2160 };
const ImgInfo FileManager::imgInfo2K_ = { Size(1920, 1080), 1920 * 1080 * sizeof(ushort), 1920 * 1080 };

int FileManager::convertTo16BitImg(const char *filename, ushort*& data, Mat& img) {
	double maxVal = 0;
	double minVal = 0;
	int ret = 0;

	log_i("Loading %s to 16bits", filename);
	if (checkFileType(filename) == 0) {
		img = imread(filename, IMREAD_ANYDEPTH | IMREAD_ANYCOLOR);
		if(img.empty()) {
			log_i("can't load %s", filename);
			return -1;
		}
	} else if (strstr(filename, ".raw")) {
        ret = rawToImg(filename, data, img);
		if(ret < 0) {
			log_i("can't load %s", filename);
			return ret;
		}
	} else if (strstr(filename, ".fits")) {
        ret = fitsToImg(filename, data, img);
		if(ret < 0) {
			log_i("can't load %s", filename);
			return ret;
		}
	}

	minMaxLoc(img, &minVal, &maxVal);
	if (maxVal <= 255)
		img.convertTo(img, CV_16UC1, 256, 0);		// 8位拉伸到16位
	else if (maxVal <= 1023)
		img.convertTo(img, CV_16UC1, 64, 0);		// 10位拉伸到16位
	else if (maxVal <= 4095)
		img.convertTo(img, CV_16UC1, 16, 0);		// 12位拉伸到16位
	else if (maxVal <= 16383)
		img.convertTo(img, CV_16UC1, 4, 0);			// 14位拉伸到16位

	return 0;
}

int FileManager::rawToImg(const char *filename, ushort*& data, Mat& img) {
	log_i("Loading %s to img", filename);

	FILE* fp = fopen(filename, "rb+");
	if(fp == NULL) {
		log_i("can't load %s", filename);
		return -1;
	}
	
	if(checkBinning(filename) < 0)
		return -1;

	data = (ushort*)calloc(imgInfo_.imgPix, sizeof(ushort));
	fread(data, sizeof(ushort), imgInfo_.imgPix, fp);
	img = Mat(imgInfo_.imgSize.height, imgInfo_.imgSize.width, CV_16UC1, data);
	fclose(fp);

	return 0;
}

int FileManager::fitsToImg(const char *filename, ushort*& data, Mat& img) {
	time_checker timer;

	// timer.start();

	fitsfile* fptr;
	int status = 0, naxis = 2, nfound, anynull;
	long naxes[2], fpixel = 1, npixels;
	ushort nullval;

	// log_i("Loading %s to img", filename);
	if (fits_open_file (&fptr, filename, READONLY, &status)) {
		fits_report_error(stderr, status);
		return -1;
	}

	// timer.stop();
	// timer.show_distance("Update stage ---- fits_open_file cost");

	if (fits_get_img_size(fptr, naxis, naxes, &status)) {
		fits_close_file(fptr, &status);
		fits_report_error(stderr, status);
		return -1;
	}

	// timer.start();

	npixels  = naxes[0] * naxes[1];         /* number of pixels in the image */
	nullval  = 0;                			/* don't check for null values in the image */

	data = (ushort*)calloc(npixels, sizeof(ushort));

	if (fits_read_img(fptr, TUSHORT, fpixel, npixels, &nullval, data, &anynull, &status)) {
		xfree(data);
		fits_close_file(fptr, &status);
		fits_report_error(stderr, status);
		return -1;
	}

	img = Mat(naxes[1], naxes[0], CV_16UC1, data);

	fits_close_file(fptr, &status);
	fits_report_error(stderr, status);

	// timer.stop();
	// timer.show_distance("Update stage ---- read fits file");

	return 0;
}

int FileManager::checkFileType(const char *filename) const {
    const static std::string types[] = {".bmp", ".png", ".jpg", ".tif"};
	const std::string filePath = filename;

    for (const std::string& type : types) {
        if (filePath.find(type) != std::string::npos) {
            return 0;
        }
    }

    return -1;
}

int FileManager::getFileSize(const char* filename, uintmax_t& fileSize) const { 
	fs::path filePath(filename);

	if (!fs::exists(filePath)) {
		log_e("FILE ERROR: %s not exist", filename);
		return -1;
	}

	fileSize = fs::file_size(filePath);

	return 0;
}

int FileManager::checkBinning(const char *filename) {
	uintmax_t fileSize;
	int ret = getFileSize(filename, fileSize);

	if (ret < 0)
		return ret;

	if (fileSize == imgInfo2K_.imgByte) 
		imgInfo_ = imgInfo2K_;
	else if (fileSize == imgInfo4K_.imgByte)
		imgInfo_ = imgInfo4K_;
	else
		return -1;

	return 0;
}

int FileManager::convertToRaw(const Mat& img, const char *filename) const {
	int nelements = img.rows * img.cols;

	log_i("%s", filename);

	if(img.empty())
		return -1;

	FILE *file = fopen(filename, "wb"); 							// open file in binary write mode
	fwrite((ushort*)img.data, sizeof(ushort), nelements, file); 	// write image data to file
	fclose(file);

	return 0;
}

/**
 * @brief 
 *      <static function> transpose a 3D image array to CHW format
 * @param dst_data 
 *      address of new buffer
 * @param src_data 
 *      address of original buffer
 * @param h 
 *      height of original image
 * @param w 
 *      width of original image
 * @param c 
 *      channel number of original image
 */
void fast_transpose_201(unsigned short* dst_data, const unsigned short* src_data, int h, int w, int c) {
	for (int i = 0; i < h; i++) {
		for (int j = 0; j < w; j++) {
			for (int k = 0; k < c; k++) {
				*(dst_data + k * h * w + i * w + j) = *(src_data + i * w * c + j * c + k);
			}
		}
	}
}

int FileManager::convertToFits(float exp, int gain, int binning, float ra, float dec, const char *object, const char *dateOBS, const Mat& imgRaw, const char *filename) const {
	time_checker timer;

	// timer.start();

	if(imgRaw.empty())
		return -1;
		
	fitsfile *fptr;
	char card[FLEN_CARD];
	int status = 0, ii = 0, jj = 0, nkeys = 0;
	long fpixel = 1, nelements = imgRaw.cols * imgRaw.rows * imgRaw.channels();

	char stringValue[FLEN_CARD] = {0};

	char filenameCopy[2048] = "!";			// 每次重写fits文件
	strcpy(filenameCopy + 1, filename);
	char* pPlace = strrchr(filenameCopy,'.');
	strcpy(pPlace,".fits");

	fits_create_file(&fptr, filenameCopy, &status);

	if (imgRaw.channels() == 1) {
		long naxis = 2;
		long naxes[2] = {imgRaw.cols, imgRaw.rows};
		fits_create_img(fptr, USHORT_IMG, naxis, naxes, &status);
	} else {
		long naxis = 3;
		long naxes[3] = {imgRaw.cols, imgRaw.rows, imgRaw.channels()};
		fits_create_img(fptr, USHORT_IMG, naxis, naxes, &status);
	}

	fits_delete_key(fptr, "COMMENT", &status);
	fits_delete_key(fptr, "COMMENT", &status);

	fits_update_key(fptr, TSTRING, "DATE-OBS", (void*)dateOBS, "Time end of exposure", &status);

	fits_update_key(fptr, TFLOAT, "EXPTIME", &exp, "Exposure Time", &status);
	fits_write_key_unit(fptr, "EXPTIME", "s", &status);
	fits_update_key(fptr, TINT, "GAIN", &gain, "Gain", &status);
	if(binning) 
		strcpy(stringValue, "2*2");
	else
		strcpy(stringValue, "1*1");
	fits_update_key(fptr, TSTRING, "BINNING", stringValue, "n*n pixels add together", &status);

	ra *= 15.0;//小时转度，Siril和PixInsight处理的是度
	fits_update_key(fptr, TFLOAT, "RA", &ra, "Observe object RA coordinate (J2000)", &status);
	fits_write_key_unit(fptr, "RA", "deg", &status);
	fits_update_key(fptr, TFLOAT, "DEC", &dec, "Observe object DEC coordinate (J2000)", &status);
	fits_write_key_unit(fptr, "DEC", "deg", &status);
	
	fits_update_key(fptr, TSTRING, "OBJECT", (void*)object, "Observe object", &status);
	strcpy(stringValue, "GBRG");
	fits_update_key(fptr, TSTRING, "BAYERPAT", stringValue, "Bayer pattern", &status);
	strcpy(stringValue, "DWARFII");
	fits_update_key(fptr, TSTRING, "TELESCOP", stringValue, "DWARF II telescope", &status);
	fits_update_key(fptr, TSTRING, "INSTRUME", stringValue, "DWARF II instrument", &status);
	strcpy(stringValue, "DWARFLAB");
	fits_update_key(fptr, TSTRING, "ORIGIN", stringValue, "DWARF II's creator", &status);

	if (imgRaw.channels() == 1) {
		fits_write_img(fptr, TUSHORT, fpixel, nelements, imgRaw.data, &status);
	} else {
		Mat img_src_rgb;
		cvtColor(imgRaw, img_src_rgb, COLOR_BGR2RGB);
		Mat img_transpose(imgRaw.size(), CV_16UC3);
		fast_transpose_201((ushort*)img_transpose.data, (ushort*)img_src_rgb.data, imgRaw.rows, imgRaw.cols, imgRaw.channels());// fits的彩图是按CHW顺序存的，每个通道一幅图
	
		fits_write_img(fptr, TUSHORT, fpixel, nelements, img_transpose.data, &status);
	}

	// fits_get_hdrspace (fptr, &nkeys, NULL, &status);

	// for (ii = 1; ii <= nkeys; ++ii)
	// {
	// 	fits_read_record (fptr, ii, card, &status);
	// 	log_i ("%s", card);
	// }

	// log_i ("END");

	fits_close_file(fptr, &status);
	fits_report_error(stderr, status);

	// timer.stop();
	// timer.show_distance("Update stage ---- write fits cost");

	return status;
}

int FileManager::convertToFits(const Mat& imgSrc, const char *filename) const {
	if(imgSrc.empty())
		return -1;

	Mat imgGray;
	cvtColor(imgSrc, imgGray, COLOR_RGB2GRAY);

	fitsfile *fptr;
	char card[FLEN_CARD];
	int status = 0, ii = 0, jj = 0, nkeys = 0;
	long fpixel = 1, naxis = 2, nelements = imgGray.cols * imgGray.rows;
	long naxes[2] = {imgGray.cols, imgGray.rows};

	char filenameCopy[2048] = "!";			// 每次重写fits文件
	strcpy(filenameCopy + 1, filename);
	char* pPlace = strrchr(filenameCopy,'.');
	strcpy(pPlace,".fits");
	log_i("%s", filenameCopy);

	fits_create_file(&fptr, filenameCopy, &status);

	fits_create_img(fptr, BYTE_IMG, naxis, naxes, &status);

	fits_delete_key(fptr, "COMMENT", &status);
	fits_delete_key(fptr, "COMMENT", &status);

	fits_write_img(fptr, TBYTE, fpixel, nelements, imgGray.data, &status);

	fits_close_file(fptr, &status);
	fits_report_error(stderr, status);

	log_i("fits status:%d", status);

	return status;
}

int FileManager::convertToColorFits(const Mat& imgSrc, const char *filename) const {
	if(imgSrc.empty())
		return -1;

	fitsfile *fptr;
	char card[FLEN_CARD];
	int status = 0, ii = 0, jj = 0, nkeys = 0;
	long fpixel = 1, naxis = 3, nelements = imgSrc.cols * imgSrc.rows * imgSrc.channels();
	long naxes[3] = {imgSrc.cols, imgSrc.rows, imgSrc.channels()};

	char filenameCopy[2048] = "!";			// 每次重写fits文件
	strcpy(filenameCopy + 1, filename);
	char* pPlace = strrchr(filenameCopy,'.');
	strcpy(pPlace,".fits");
	log_i("%s", filenameCopy);

	Mat img_src_rgb;
	cvtColor(imgSrc, img_src_rgb, COLOR_BGR2RGB);
	Mat img_transpose(imgSrc.size(), CV_16UC3);
	fast_transpose_201((ushort*)img_transpose.data, (ushort*)img_src_rgb.data, imgSrc.rows, imgSrc.cols, imgSrc.channels());// fits的彩图是按CHW顺序存的，每个通道一幅图

	fits_create_file(&fptr, filenameCopy, &status);

	fits_create_img(fptr, USHORT_IMG, naxis, naxes, &status);
	fits_write_img(fptr, TUSHORT, fpixel, nelements, img_transpose.data, &status);

	fits_delete_key(fptr, "COMMENT", &status);
	fits_delete_key(fptr, "COMMENT", &status);

	fits_close_file(fptr, &status);
	fits_report_error(stderr, status);

	log_i("fits status:%d", status);

	return status;
}

int FileManager::convertToFitsFromFile(const char *filename) const {
	Mat imgSrc = imread(filename, IMREAD_ANYDEPTH | IMREAD_ANYCOLOR);
	if(imgSrc.empty()) {
		log_i("can't load %s", filename);
		return -1;
	}

	Mat imgGray;
	cvtColor(imgSrc, imgGray, COLOR_RGB2GRAY);

	fitsfile *fptr;
	char card[FLEN_CARD];
	int status = 0, ii = 0, jj = 0, nkeys = 0;
	long fpixel = 1, naxis = 2, nelements = imgGray.cols * imgGray.rows;
	long naxes[2] = {imgGray.cols, imgGray.rows};

	char filenameCopy[2048] = "!";			// 每次重写fits文件
	strcpy(filenameCopy + 1, filename);
	char* pPlace = strrchr(filenameCopy,'.');
	strcpy(pPlace,".fits");
	log_i("%s", filenameCopy);

	fits_create_file(&fptr, filenameCopy, &status);

	fits_create_img(fptr, BYTE_IMG, naxis, naxes, &status);

	fits_delete_key(fptr, "COMMENT", &status);
	fits_delete_key(fptr, "COMMENT", &status);

	fits_write_img(fptr, TBYTE, fpixel, nelements, imgGray.data, &status);

	fits_close_file(fptr, &status);
	fits_report_error(stderr, status);

	return status;
}

